#
# Copyright (c) 2025 Huawei Technologies Co., Ltd.
# openUBMC is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#         http://license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.
#

"""
嵌套函数类型推断工具
"""

from typing import Optional

from mcpy.codegen.core.context import CodeGenContext
from mcpy.codegen.core.generator import get_global_generator
from mcpy.ir.statements.assignment import IRAssign
from mcpy.ir.statements.control import IRIf
from mcpy.ir.expressions.basic import IRName
from mcpy.ir.expressions.call import IRCall
from mcpy.type_system import is_protocol_type


def get_await_result_type(target_node) -> Optional[str]:
    """获取 await 结果变量的 C++ 类型

    Args:
        target_node: await 赋值目标节点 (IRName)

    Returns:
        C++ 类型字符串，如果无法确定则返回 None
    """
    if not target_node:
        return None

    # 优先使用已推断的 cpp_type
    cpp_type = target_node.get_extension_param("cpp_type")
    if cpp_type:
        return cpp_type

    # 否则从 inferred_type 转换
    inferred_type = target_node.get_extension_param("inferred_type")
    if inferred_type:
        type_str = str(inferred_type)
        # 简单类型映射
        type_map = {
            "int": "int",
            "str": "std::string",
            "bool": "bool",
            "float": "double",
        }
        return type_map.get(type_str, None)

    return None


def get_local_var_type(
    statements: list, var_name: str, context: CodeGenContext
) -> Optional[str]:
    """递归查找局部变量的类型

    Args:
        statements: 语句列表
        var_name: 变量名
        context: 代码生成上下文

    Returns:
        C++ 类型字符串，如果找不到返回 None
    """
    generator = get_global_generator()

    def search_in_stmts(stmts):
        for stmt in stmts:
            if isinstance(stmt, IRAssign):
                target = stmt.get_target_expr()
                if (
                    target
                    and isinstance(target, IRName)
                    and getattr(target, "id", None) == var_name
                ):
                    # 找到定义，获取类型
                    # 优先使用已经转换好的 cpp_type（类型替换阶段的结果）
                    cpp_type = target.get_extension_param("cpp_type")
                    if not cpp_type:
                        # 如果没有 cpp_type，尝试从 inferred_type 转换
                        target_inferred_type = target.get_extension_param("inferred_type")
                        if target_inferred_type:
                            from mcpy.type_system import get_cpp_type
                            cpp_type = get_cpp_type(target_inferred_type, target)
                        else:
                            cpp_type = context.get_cpp_type(stmt.value)

                    # 检查是否是 Protocol 类型（映射到 mc::variant）
                    # 如果是且有类型推导信息，使用 decltype
                    if cpp_type == "mc::variant":
                        inferred_type = target.get_extension_param("inferred_type")
                        if inferred_type and is_protocol_type(inferred_type):
                            # 对于 Protocol 类型，生成基于 decltype 的类型声明
                            # 使用 std::declval 避免依赖变量声明顺序
                            # 例如：iter(list[int]) -> std::decay_t<decltype(mcpy::iter_wrapper(std::declval<mc::array<int>>()))>
                            try:
                                # 检查是否是 iter() 调用
                                if isinstance(stmt.value, IRCall):
                                    func = stmt.value.func
                                    if (
                                        getattr(func, "id", None) == "iter"
                                        and stmt.value.args
                                    ):
                                        # 获取参数的类型
                                        arg_cpp_type = context.get_cpp_type(
                                            stmt.value.args[0]
                                        )
                                        if (
                                            arg_cpp_type
                                            and arg_cpp_type != "mc::variant"
                                        ):
                                            # 使用 std::declval 生成类型声明
                                            cpp_type = f"std::decay_t<decltype(mcpy::iter_wrapper(std::declval<{arg_cpp_type}>()))>"
                            except Exception:
                                # 如果生成失败，fallback 到 mc::variant
                                pass

                    return cpp_type

            elif isinstance(stmt, IRIf):
                # 递归查找 if 分支
                if stmt.then_body:
                    result = search_in_stmts(stmt.then_body.statements)
                    if result:
                        return result

                if stmt.else_body:
                    result = search_in_stmts(stmt.else_body.statements)
                    if result:
                        return result
            
            elif hasattr(stmt, 'body') and stmt.body:
                # 递归查找循环体（IRWhile, IRForEach, IRForRange 等）
                if hasattr(stmt.body, 'statements'):
                    result = search_in_stmts(stmt.body.statements)
                    if result:
                        return result
        return None

    return search_in_stmts(statements)
